#include "core/task.h"
#include "tools/klib.h"
#include "tools/log.h"
#include "cpu/cpu.h"
#include "cpu/irq.h"
#include "cpu/mmu.h"
#include "core/memory.h"
#include "comm/cpu_instr.h"
#include "os_cfg.h"

static task_manager_t task_manager;
static uint32_t idle_stack[IDLE_TASK_STACK_SIZE];   // 空闲进程的栈

static int tss_init(task_t *task, uint32_t entry, uint32_t esp)
{
    int tss_sel = gdt_alloc_desc();
    if (tss_sel < 0) {
        log_printf("alloc tss failed.\n");
        return -1;
    }

    segment_desc_set(tss_sel, (uint32_t)&task->tss, sizeof(tss_t) - 1, 
                    SEG_P_PRESENT | SEG_DPL0 | SEG_TYPE_TSS);

    kernel_memset(task, 0, sizeof(task_t));

    int code_sel = task_manager.app_code_sel | SEL_CPL3;
    int data_sel = task_manager.app_data_sel | SEL_CPL3;

    task->tss.eip = entry;
    task->tss.esp = task->tss.esp0 = esp;
    task->tss.ss = data_sel;
    task->tss.ss0 = KERNEL_SELECTOR_DS;
    task->tss.ds = task->tss.es = task->tss.fs = task->tss.gs = data_sel;
    task->tss.cs = code_sel;
    task->tss.eflags = EFLAGS_DEFAULT | EFLAGS_IF;

    uint32_t page_dir = memory_create_uvm();
    if (page_dir == 0) {
        gdt_free_sel(tss_sel);
        return -1;
    }
    task->tss.cr3 = page_dir;

    task->tss_sel = tss_sel;
    return 0;
}

int task_init(task_t *task, uint32_t entry, uint32_t esp, const char *pname)
{
    ASSERT(task != 0);

    tss_init(task, entry, esp);

    kernel_strncpy(task->name, pname, TASK_NAME_SIZE);
    task->state = TASK_CREATE;
    task->sleep_ticks = 0;
    task->time_ticks = TASK_TIME_SLICE_DEFAULT;
    task->slice_ticks = task->time_ticks;
    list_node_init(&task->run_node);
    list_node_init(&task->all_node);
    list_node_init(&task->wait_node);
   
    //操作任务队列需要临界区保护
    irq_state_t state = irq_enter_protection();
    task_set_ready(task);
    list_insert_last(&task_manager.task_list, &task->all_node);
    irq_leave_protection(state);
    return 0;
}

void task_switch(task_t *from, task_t *to)
{
    switch_to_tss(to->tss_sel);
}

static void idle_task_entry(void)
{
    while(1)
        hlt();
}

void task_manager_init(void)
{
    // 定义用户程序数据段
    int app_data_sel = gdt_alloc_desc();
    segment_desc_set(app_data_sel, 0x00000000, 0xFFFFFFFF, 
                    SEG_P_PRESENT | SEG_DPL3 | SEG_S_NORMAL | SEG_TYPE_DATA 
                    | SEG_TYPE_RW | SEG_D);
    task_manager.app_data_sel = app_data_sel;

    // 定义用户程序代码段
    int app_code_sel = gdt_alloc_desc();
    segment_desc_set(app_code_sel, 0x00000000, 0xFFFFFFFF, 
                    SEG_P_PRESENT | SEG_DPL3 | SEG_S_NORMAL | SEG_TYPE_CODE 
                    | SEG_D);
    task_manager.app_code_sel = app_code_sel;

    list_init(&task_manager.ready_list);
    list_init(&task_manager.task_list);
    list_init(&task_manager.sleep_list);
    task_manager.curr_task = NULL;

    task_init(&task_manager.idle_task, 
            (uint32_t)idle_task_entry, 
            (uint32_t)(idle_stack + IDLE_TASK_STACK_SIZE), 
            "idle task");
}

void first_task_init(void)
{
    void first_task_entry(void);
    extern uint8_t s_first_task[], e_first_task[];

    uint32_t copy_size = (uint32_t)(e_first_task - s_first_task);
    uint32_t alloc_size = 10 * MEM_PAGE_SIZE;
    ASSERT(copy_size < alloc_size);
    
    uint32_t first_start = (uint32_t)first_task_entry;

    task_init(&task_manager.first_task, first_start, 0, "first_task");
    ltr(task_manager.first_task.tss_sel);
    task_manager.curr_task = &task_manager.first_task; 

    mmu_set_page_dir(task_manager.first_task.tss.cr3);

    memory_alloc_page_for(first_start, alloc_size, PTE_P | PTE_W);
    kernel_memcpy((void *)first_start, s_first_task, copy_size);
}

task_t * get_first_task(void)
{
    return &task_manager.first_task;
}

task_t * get_current_task(void)
{
    return task_manager.curr_task;
}

void task_set_ready(task_t *task)
{
    // 空闲进程不进入就绪队列
    if (task == &task_manager.idle_task)
        return;

    list_insert_last(&task_manager.ready_list, &task->run_node);
    task->state = TASK_READY;
}

void task_set_block(task_t *task)
{
    if (task == &task_manager.idle_task)
        return;
    list_remove(&task_manager.ready_list, &task->run_node);
}

task_t * task_next_run(void)
{
    if (list_count(&task_manager.ready_list) == 0) {
        return &task_manager.idle_task;
    }
    list_node_t *first_node = list_first(&task_manager.ready_list);
    return field_struct_addr(first_node, task_t, run_node);
}

int sys_sched_yield(void)
{
    //操作任务队列需要临界区保护
    irq_state_t state = irq_enter_protection();
    if (list_count(&task_manager.ready_list) > 1) {
        task_t *curr_task = get_current_task();
        task_set_block(curr_task);
        task_set_ready(curr_task);
        task_dispatch();
    }
    irq_leave_protection(state);
    return 0;
}

void task_dispatch(void)
{
    //操作任务队列需要临界区保护
    irq_state_t state = irq_enter_protection();

    task_t *from = get_current_task();
    task_t *to = task_next_run();
    if (from != to) {
        task_manager.curr_task = to;
        to->state = TASK_RUNNING;
        task_switch(from, to);
    }

    irq_leave_protection(state);
}

void task_time_tick(void)
{
    //操作任务队列需要临界区保护
    irq_state_t state = irq_enter_protection();

    // 如果当前任务时间片消耗完，则切换任务
    task_t *curr_task = get_current_task();
    if (curr_task->slice_ticks-- <= 0) {
        curr_task->slice_ticks = curr_task->time_ticks;
        task_set_block(curr_task);
        task_set_ready(curr_task);
        task_dispatch();
    }
    
    // 遍历休眠队列，唤醒休眠时间片到期的进程
    list_node_t *curr = list_first(&task_manager.sleep_list);
    while (curr) {
        list_node_t *next = list_node_next(curr);
        task_t *task = field_struct_addr(curr, task_t, run_node);
        if (--task->sleep_ticks == 0) {
            task_set_wakeup(task);
            task_set_ready(task);
        }
        curr = next;
    }

    task_dispatch();

    irq_leave_protection(state);
}

void sys_sleep(uint32_t ms)
{
    irq_state_t state = irq_enter_protection();

    task_set_block(task_manager.curr_task);
    // 向上取整得到休眠的时钟节拍数
    task_set_sleep(task_manager.curr_task, (ms + OS_TICK_MS - 1) / OS_TICK_MS);
    task_dispatch();

    irq_leave_protection(state);
}

void task_set_sleep(task_t *task, uint32_t ticks)
{
    if (ticks == 0)
        return;

    task->sleep_ticks = ticks;
    task->state = TASK_SLEEP;
    list_insert_last(&task_manager.sleep_list, &task->run_node);
}

void task_set_wakeup(task_t *task)
{
    list_remove(&task_manager.sleep_list, &task->run_node);
}
